package com.yitiao.ratelimit;

public class TokenBucket {
    private int capacity; // 桶的容量
    private int rate; // 令牌放入速度
    private int tokens; // 当前令牌数量
    private long timestamp; // 时间戳

    public TokenBucket(int capacity, int rate) {
        this.capacity = capacity;
        this.rate = rate;
        this.tokens = capacity;
        this.timestamp = System.currentTimeMillis();
    }

    public synchronized boolean acquire(int num) {
        refill();
        if (tokens >= num) {
            tokens -= num;
            return true;
        } else {
            return false;
        }
    }

    private void refill() {
        long now = System.currentTimeMillis();
        int addTokens = (int) ((now - timestamp) / 1000 * rate);
        if (addTokens > 0) {
            tokens = Math.min(capacity, tokens + addTokens);
            timestamp = now;
        }
    }
}

